"""
季节指数
"""
import numpy as np
import pandas as pd
from sklearn.metrics import mean_squared_error


def ma(series, window):
    """
    移动平均趋势剔除法
    :param series:
    :param window: 移动窗口
    :return: 调整后的季节指数
    """
    # 中心化移动均值
    data_ma = series.rolling(window=window).mean()
    data_ma_cma = data_ma.rolling(window=2).mean()
    # 季节比率
    df = pd.DataFrame([series, data_ma, data_ma_cma]).T
    df.columns = ['raw', 'ma', 'cma']
    df['window_rate'] = df['raw'] / df['cma']

    range_num = int(np.ceil(len(series) / window))
    # 季节比率-分组
    window_rate_arr = \
        [df['window_rate'].iloc[i * window: (i + 1) * window].tolist()
         for i in range(1, range_num)]
    # 季节比率-window 均值
    window_rate_mean = np.mean(window_rate_arr, axis=0)
    # 季节指数调整
    window_mean = window_rate_mean.mean()
    window_rate_mean_adj = [num / window_mean for num in window_rate_mean]
    return window_rate_mean_adj


def simple_seasonal_index(series, window):
    """
    简单季节指数法
    :param series:
    :param window:
    :return:
      季节指数 seasonal_index
      预测参数a data_mean,
      预测参数b weight
      观察周期 参数T observate_period
    """
    data_set_length = len(series)
    # 各年同月/季观察值的均值(A)
    range_num = int(np.ceil(data_set_length / window))
    window_rate_arr = \
        [series[i * window: (i + 1) * window].tolist()
         for i in range(1, range_num)]

    # 数据填充
    while window_rate_arr and len(window_rate_arr[-1]) < window:
        window_rate_arr[-1].append(np.nan)

    window_rate_mean = np.nanmean(window_rate_arr, axis=0)
    if isinstance(window_rate_mean, np.float):
        return None, None, None, None
    # 历年所有月份/季度的均值(B)
    data_mean = series.mean()
    # 季节指数
    seasonal_index = [num / data_mean for num in window_rate_mean]

    # 预测参数
    observate_period = []
    if data_set_length % 2 == 0:
        for num in range(int(data_set_length / 2)):
            observate_period.append(-1 - 2 * num)
            observate_period.append(1 + 2 * num)
    else:
        observate_period.append(0)
        for num in range(1, int((data_set_length + 1) / 2)):
            observate_period.append(-num)
            observate_period.append(num)
    observate_period.sort()

    # 分子
    weight_molecular_arr = \
        [num * observate_period[idx]
         for idx, num in enumerate(series)]
    # 分母
    weight_denominator = np.power(observate_period, 2).sum()
    weight = sum(weight_molecular_arr) / weight_denominator

    return seasonal_index, data_mean, weight, observate_period


def seasonal_index_main(series, window):
    """
    季节指数
    :param series:
    :param window: 周期
    :return:
    """
    if isinstance(series, pd.Series):
        data_set = series.to_numpy()

    if not isinstance(series, np.ndarray):
        data_set = np.array(series)
    else:
        data_set = series

    rms_loss = []

    for predict_period in range(6, 10):
        train = data_set[: len(data_set) - predict_period]
        test = data_set[len(data_set) - predict_period:]

        seasonal_index, intercep, coef, observate_period = \
            simple_seasonal_index(series=train, window=window)

        if seasonal_index is None:
            print('参数异常！')
            return

        # predict
        observate_period_step = observate_period[-1] - observate_period[-2]
        pred_arr = []
        # 训练集的最后一个周期对应的月/季度值
        latest_period_idx = len(data_set) % window
        for idx in range(len(test)):
            # 月/季度值idx
            period_idx = (latest_period_idx + idx) % window
            # 预测周期T值
            predict_peroid = \
                (idx + 1) * observate_period_step + observate_period[-1]
            # 预测值
            mid_pred = intercep + coef * predict_peroid

            pred_arr.append(mid_pred * seasonal_index[period_idx])

        # 评估
        rms_loss.append(np.sqrt(mean_squared_error(test, pred_arr)))

    print('rms: {}'.format(np.mean(rms_loss)))
